#!/usr/bin/env python

import sys, os

Lmap = { "s":0, "p":1, "d":2 }
# Dictionary for number of cartesian components in a shell
lm = { "s":1, "p":3, "d":6, "f":10 }
# Dictionary for number of spherical components in a shell
#lm = { "s":1, "p":3, "d":5, "f":7 }

#-----------------------------------------------------------------------------

class Atoms:

   def __init__(self,natoms):
      self.charge = []
      self.ncont = []
      self.nprim = []
      self.shells = []
      self.xyz = []
      self.natoms = natoms
      self.nshells = 0

   def put_atoms_and_shells(self,lines):
      lines.pop()  # blank
      nshells = 0
      for k in range(self.natoms):
         line = lines.pop()
         foo = line.split()
         bar = foo[5].rstrip("]").split("|")[1]
         shells = {}
         for u in ["s","p","d","f","g","h","i"]:
            if len(bar) > 1:
               bar = bar.split(u)
               shells[u] = eval(bar[0])
               nshells += eval(bar[0])
               bar = bar[1]
         self.charge.append(eval(foo[2]))
         self.nprim.append(eval(foo[3]))
         self.ncont.append(eval(foo[4]))
         self.shells.append(shells)
      self.nshells = nshells

   def put_xyz(self,lines):
      for k in range(self.natoms):
         lines.pop()  # blank
         x = lines.pop().split()[3]
         y = lines.pop().split()[2]
         z = lines.pop().split()[2]
         self.xyz.append((x,y,z))

#-----------------------------------------------------------------------------

class Shells:

   def __init__(self,atoms,lines):
      self.atom_id = []
      self.L = []
      self.exp_id = []
      self.nprim = []
      self.basis = []

      lines.pop()  # blank
      lines.pop()  # blank

      for i in range(atoms.natoms):
         for u in ["s","p","d","f","g","h","i"]:
            if u in atoms.shells[i]:
               for j in range(atoms.shells[i][u]):
                  line = lines.pop()
                  if u in lm:
                     for k in range(lm[u]-1): lines.pop()  # skip
                  else:
                     print " Input # of Cartesian components for L =", u
                  self.atom_id.append(i)
                  self.L.append(u)
                  self.exp_id.append((j,line.split()[3:]))
                  self.nprim.append(len(line.split()[3:]))

#-----------------------------------------------------------------------------

class Basis:

   def __init__(self,atoms,lines):

      self.exp = []
      self.coef = []
      self.nprim = 0

      lines.pop()  # blank
      lines.pop()  # blank
      lines.pop()  # blank
      while len(lines) > 0:
         line = lines.pop()
         if line.find("Copy of input to READIN") > 0: break
         if len(line) > 1:
            line = line[13:].split()
            self.exp.append(line[1])
            self.coef.append(line[2:])
      self.nprim = len(self.exp)

#-----------------------------------------------------------------------------

def get_atoms_and_basis_funcs(dalton):

   print " Reading atoms and basis data from DALTON output file..."

   lines = dalton.readlines()
   lines.reverse()
   while len(lines) > 0:
      line = lines.pop()
      if line.find("Total number of atoms:") > 0:
         natoms = eval(line.split()[4])
         break

   atoms = Atoms(natoms)   

   while len(lines) > 0:
      line = lines.pop()
      if line.find("label    atoms   charge   prim    cont     basis") > 0:
         atoms.put_atoms_and_shells(lines)
      if line.find("Total number of coordinates") > 0:
         atoms.put_xyz(lines)
         break
   print " Number of atoms = ", atoms.natoms
   print " Charges = ", atoms.charge
   print " Number of contracted funcs = ", atoms.ncont
   print " Number of primitive funcs = ", atoms.nprim
   print " Shell structure = ", atoms.shells
   if (lm["d"] == 6): print " Assuming Cartesian d functions!"
   #print atoms.xyz

   while len(lines) > 0:
      line = lines.pop()
      if line.find("Contracted Orbitals") > 0:
         shells = Shells(atoms,lines)   
         break

   while len(lines) > 0:
      line = lines.pop()
      if line.find("Orbital exponents and normalized contraction") > 0:
         basis = Basis(atoms,lines)   
         break

   return atoms, shells, basis

#-----------------------------------------------------------------------------

def process_for_qcim_format(atoms,shells,basis):

   # Expand exponents and coeffs into one array for each shell
   # Change array indices +1 for Fortran
   shells.basis = []
   exp_id = 0
   for sh in range(len(shells.L)):
      tmp = []
      subsh = shells.exp_id[sh][0]
      for prim in shells.exp_id[sh][1]:
         prim = eval(prim) -1
         tmp.append((basis.exp[prim],basis.coef[prim][subsh]))

      shells.basis.append(tmp)
      shells.exp_id[sh] = exp_id + 1
      exp_id += shells.nprim[sh]
      shells.atom_id[sh] += 1

   #print
   #for sh in range(len(shells.L)):
   #   print shells.atom_id[sh], shells.L[sh], shells.exp_id[sh]
   #   for j in shells.basis[sh]:
   #      print j

   # Convert "s,p,d..." string to integer value "0,1,2..."
   for sh in range(len(shells.L)):
      shells.L[sh] = Lmap[shells.L[sh]]

#-----------------------------------------------------------------------------

def write_fortran_basis_file(atoms,shells,basis):

   print " Writing Fortran interface file..."

   # Nuclei co-ordinates and charges for FMM interface
   qcim = open("nuclei.qcim","w")
   qcim.write(str(atoms.natoms)+"\n")
   for i in range(len(atoms.xyz)):
      qcim.write(str(atoms.xyz[i][0])+"\n")
      qcim.write(str(atoms.xyz[i][1])+"\n")
      qcim.write(str(atoms.xyz[i][2])+"\n")
      qcim.write(str(atoms.charge[i])+"\n")

   # Basis set information for MD1 interface
   qcim = open("basis.qcim","w")

   qcim.write(str(atoms.natoms)+"\n")
   qcim.write(str(atoms.nshells)+"\n")
   qcim.write(str(sum(atoms.ncont))+"\n")
   qcim.write(str(sum(shells.nprim))+"\n")

   for i in atoms.xyz:
      qcim.write(str(i[0])+"\n")
      qcim.write(str(i[1])+"\n")
      qcim.write(str(i[2])+"\n")

   for i in range(len(shells.L)):
      qcim.write(str(shells.atom_id[i])+"\n")
      qcim.write(str(shells.L[i])+"\n")
      qcim.write(str(shells.exp_id[i])+"\n")
      qcim.write(str(shells.nprim[i])+"\n")
      for bas_coef in shells.basis[i]:
         qcim.write(str(bas_coef[0])+"\n")
         qcim.write(str(bas_coef[1])+"\n")

   qcim.close()

#-----------------------------------------------------------------------------

if __name__ == '__main__':

   print " -----------------------------------"
   print " |   QCIM wrapper for DALTON 2.0   |"
   print " -----------------------------------"
   
   name = str(sys.argv[1])
   # Run standard DALTON first
#   print " Running Dalton 2.0 to generate reference output."
#   os.system("/home/mark/science/apps/dalton/dalton-2.0-cam/bin/dalton -D "+name)
#   os.system("cp "+name+".out "+name+".ref.out")
#   print " -----------------------------------"
   
   # Parse Dalton output file to generate QCIM interface files
   dalton = open(name+".ref.out","r")
   atoms, shells, basis = get_atoms_and_basis_funcs(dalton)
   dalton.close()

   # Reorganize basis data for QCIM format
   process_for_qcim_format(atoms,shells,basis)

   # Write-out interface file for F90 QCIM library
   write_fortran_basis_file(atoms,shells,basis)
   os.system("! [ -d /tmp/dalton-qcim/"+name+" ] && mkdir /tmp/dalton-qcim/"+name)
   os.system("cp -f nuclei.qcim /tmp/dalton-qcim/"+name)
   os.system("cp -f nuclei.qcim nuclei.fmm")
   os.system("cp -f nuclei.fmm /tmp/dalton-qcim/"+name)
   os.system("cp -f basis.qcim /tmp/dalton-qcim/"+name)
   os.system("cp -f geom.qcim /tmp/dalton-qcim/"+name)

   # Now run DALTON for QCIM calculation
   print " Running Dalton 2.0 with QCIM updates..."
   os.system("/home/mark/science/apps/dalton/dalton-2.0-qcim/bin/dalton -D "+name)
   os.system("grep 'Final.*energy' "+name+"*.out*")

   # Combine output files into one
   os.system("cp /tmp/dalton-qcim/"+name+"/qcim.out .")
   os.system("cat qcim.out >> "+name+".out")
   os.system("cat geom.qcim >> "+name+".out")
   os.system("cp "+name+".out "+name+".qcim.out")

